#! /usr/bin/env python

'''
grepidemic.py - Modeling information flow through a network. Uses NodeBox to provide a gui 
for simulations. Requires NodeBox
'''

__author__  = "Chris Hill"
__version__ = "0.1"
__license__ = "GPL"

import sys, os
graph =  ximport("graph")
coreimage = ximport("coreimage")

FRAMES_PER_SECOND = 30
DAYS_PER_SECOND = 10
DEFAULT_INFECTION_RATE = .7

DYNAMIC_INFECT = True
DYNAMIC_PROB = DEFAULT_INFECTION_RATE * 100 + 1.5

DEBUG = True

filename = 'graphs/pl_25_3_60'

class InfectedGraph(object):
    """InfectedGraph models the spread of a infectious agent through a graph.
    This agent can be the flu through a social network, or re-tweeting news
    through the twitter social network.
    
    Attributes:
        g = graph
        name = name of the graph
        infected_date = most recent node infection
        prob_infection = probability of an infection node spreading to it's neighbor
        frame = timeshot of the universe
        todo(cmhill)
    """
    
    def __init__(self, g, name = ""):
        self.g = g
        self.infected_nodes = []
        self.frame = 0
        self.name = name
        self.prob_infection = DEFAULT_INFECTION_RATE
        self.infection_length = 30
        
        self.infected_date = {}
        self.infected_countdown = {}
        self.chance_of_infection = {}
        
        self.g.events.popup = True
        self.g.events.click = self.infectNode
        
    
    def header(self, txt, x = 50, y = 75, size = 55):
        w = 200
        
        fill(180.0/255, 198.0/255, 177.0/255, 1)
        font("Courier", size)
        text(txt, x, y, outline=True)
        return textheight(txt)
        
    def infectNode(self, node):
        node.style = "important"
        self.infected_date[node.id] = str(self.frame / FRAMES_PER_SECOND)
        #node.style = "1, 0, 0.25, 0.75"
        self.infected_nodes.append(node)
        self.infected_countdown[node.id] = self.infection_length
        print '[OUTBREAK] ' + node.id + ", @ t = " + str(self.frame)
    
    def getNodeStats(self, g, node):
        info = 'Infected on day: ' + self.infected_date.get(node, 'NA')
        
        info += '\nInfection countdown: '
        info += str(self.infected_countdown.get(node, 'NA'))
        '''if self.infected_countdown.get(node, 'NA') is not 'NA':
            info += str(self.infected_countdown.get(node, 30) / DAYS_PER_SECOND)
        else:
            info += 'NA'''
            
        info += '\nChance of infection: ' + str(self.chance_of_infection.get(node, 'NA'))
        return info
        
        #return len(g.nodes)
    
    def draw(self, iterations = 10):
        
        global FRAMES_PER_SECOND
        
        # Calculate percent of getting infected next iteration
        for node in g.nodes:
            self.g.events.popup_text[node.id] = self.getNodeStats(self.g, node.id)
        
        if keydown and keycode == 49:
            self.header("Day " + str(self.frame / 10), 10,75) #10, 975)
            self.header(self.name, 10, 20, size = 20)
            self.g.draw(weighted = True)
            self.g.update(iterations)
        else:
            self.header("Day " + str(self.frame / 10), 10,75)#10, 975)
            self.header(self.name, 10, 20, size = 20)
        
            for node in g.nodes:
                infect_nodes = []
                immune_nodes = []
                
                # Grow infected nodes, and shrink immune nodes.
                if node.style is 'important':
                    node.r = node.r * 1.025
                    
                    for neighbor in node.links:
                            #neighbor.style = 'important'
                            if neighbor.style is not 'highlight': infect_nodes.append(neighbor)                    
                        
                elif node.style is 'highlight':
                    node.r = node.r * .9756
                    if node.r < 8.0:
                        node.style = 'default'
                        node.r = 10.0
                else:
                    for neighbor in node.links:
                    #neighbor.style = 'important'
                        if neighbor.style  is 'important':
                            pass
                    #self.chance_of_infection
                # Randomly infect another group of nodes
                for inode in infect_nodes:
                    if inode.style is not 'highlight' and inode.style is not 'important':
                        if self.prob_infection * 100 < random(DYNAMIC_PROB) and DYNAMIC_INFECT:
                            self.infected_date[inode.id] = str(self.frame / DAYS_PER_SECOND)
                            self.g.events.popup_text[inode.id] = self.getNodeStats(self.g, inode.id)
                            inode.style = 'important'
                            inode.r = 8
                            self.infected_nodes.append(inode)
                            self.infected_countdown[inode.id] = self.infection_length
                            print '[INFECT] ' + node.id + ' -> ' + inode.id + ' @ t = ' + str(self.frame)
    
                #for node in immune_nodes:
                #    node.style = 'default'        
                #    node.r = 8.0

            #for node in self.
    
            # Go through all infected nodes and decrease their sickness timer.
            # Also spread the illness by some chance.
            for node in self.infected_nodes:
                self.infected_countdown[node.id] = self.infected_countdown.get(node.id, FRAMES_PER_SECOND) - 1
                
                if self.infected_countdown[node.id] < 0:
                    node.style = 'highlight'
                    self.infected_nodes.remove(node)
                    #node.r = 8.0
                
            if self.frame % FRAMES_PER_SECOND == 60:
                infect_nodes = []
                immune_nodes = []
                for node in self.g.nodes:
                    if node.style is 'important':
                        for neighbor in node.links:
                            #neighbor.style = 'important'
                            if neighbor.style is not 'highlight': infect_nodes.append(neighbor)
                
                        node.style = 'highlight'
                
                    elif node.style is 'highlight':
                        immune_nodes.append(node)
                        #node.style = 'default'
        
                for node in infect_nodes:
                    if node.style is not 'highlight':
                        if self.prob_infection * 100 > random(100):
                            self.infected_date[node.id] = str(self.frame / DAYS_PER_SECOND)
                            self.g.events.popup_text[node.id] = self.getNodeStats(self.g, node.id)
                            node.style = 'important'
                            node.r = 8
                            self.infected_nodes.append(node)
        
                for node in immune_nodes:
                    node.style = 'default'        
                    node.r = 8.0
                    
            elif random(100) < 0:
                infect_nodes = []
                immune_nodes = []
                for node in self.g.nodes:
                    if node.style is 'important':
                        for neighbor in node.links:
                            #neighbor.style = 'important'
                            if neighbor.style  is not 'highlight': infect_nodes.append(neighbor)
                
                        node.style = 'highlight'
                
                    elif node.style is 'highlight':
                        immune_nodes.append(node)
                        #node.style = 'default'
        
                for node in infect_nodes:
                    if node.style is not 'highlight':
                        if self.prob_infection * 100 < random(100):
                            self.infected_date[node.id] = str(self.frame / DAYS_PER_SECOND)
                            self.g.events.popup_text[node.id] = self.getNodeStats(self.g, node.id)
                            node.style = 'important'
                            node.r = 8
                            self.infected_nodes.append(node)
                            
                            #print '[INFECT] ' + node.id + ' -> ' + inode.id

            self.g.draw(weighted = True)#, 300, 300)
            self.g.update(iterations)
            self.frame += 1
    
global ig

g = graph.create(iterations=1000, distance=1.8, depth=False)

        
file = open(filename, 'r')
for line in file:
    tuple = line.strip().split('\t')
    g.add_edge(tuple[0], tuple[1])

g.solve()

# Created the infected graph
ig = InfectedGraph(g, filename)

g.draw()
speed(30)

def draw():
    global ig
    ig.draw(10)